import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.stats as stats
from osc05.static.constants import Constants


class PERTGenerator:
    def __init__(self):
        self.params_df = pd.read_csv('sampling_parameters.csv')
        self.current_group = None
        self.samples = None
        self.results = []
        self.valid_samples = {}  # 存储通过检验的样本组

    def pert_samples(self, a, b, c, size, rng):
        """Generate PERT distribution samples"""
        mu = (a + 4 * b + c) / 6
        alpha = ((mu - a) * (2 * b - a - c)) / ((b - mu) * (c - a))
        beta = (alpha * (c - mu)) / (mu - a)
        return a + (c - a) * rng.beta(alpha, beta, size=size)

    def clip_by_row(self, samples, special_rows, a_spec, c_spec, a_std, c_std):
        """Clip samples based on row-specific bounds"""
        bounds = np.full((samples.shape[0], 2), [a_std, c_std])
        bounds[special_rows] = [a_spec, c_spec]
        return np.clip(samples, bounds[:, 0:1], bounds[:, 1:2])

    def check_pert_distribution(self, samples, a, b, c, group_id, row_idx=0, plot=True):
        """Validate PERT distribution with 4 diagnostic plots and KS test"""
        data = samples[row_idx]
        mu = (a + 4 * b + c) / 6
        alpha = ((mu - a) * (2 * b - a - c)) / ((b - mu) * (c - a))
        beta_val = (alpha * (c - mu)) / (mu - a)

        theoretical_dist = stats.beta(alpha, beta_val, loc=a, scale=c - a)
        ks_stat, p_value = stats.kstest(data, theoretical_dist.cdf)

        if plot:
            self._plot_distribution(data, a, b, c, alpha, beta_val, group_id,
                                    row_idx, theoretical_dist, ks_stat, p_value)

        return {
            'alpha': alpha,
            'beta': beta_val,
            'ks_stat': ks_stat,
            'p_value': p_value,
            'ks_pass': p_value > 0.05
        }

    def _plot_distribution(self, data, a, b, c, alpha, beta, group_id,
                           row_idx, theoretical_dist, ks_stat, p_value):
        """Generate 4-panel diagnostic plot with group info"""
        x = np.linspace(a, c, 200)
        pdf = theoretical_dist.pdf(x)
        cdf = theoretical_dist.cdf(x)

        fig, axs = plt.subplots(2, 2, figsize=(15, 12))
        fig.suptitle(
            f'Group {group_id} - Row {row_idx} PERT Distribution Validation\n'
            f'a={a}, b={b}, c={c} | α={alpha:.2f}, β={beta:.2f}\n'
            f'KS Stat: {ks_stat:.4f}, p-value: {p_value:.4f} '
            f'({"PASS" if p_value > 0.05 else "FAIL"})',
            fontsize=14, y=1.02
        )

        # 1. Histogram vs PDF
        axs[0, 0].hist(data, bins=30, density=True, alpha=0.6,
                       color='steelblue', label='Sampled Data')
        axs[0, 0].plot(x, pdf, 'r-', lw=2, label='Theoretical PDF')
        axs[0, 0].set_title('Density Comparison')
        axs[0, 0].legend()

        # 2. Q-Q Plot
        stats.probplot(data, dist=theoretical_dist, plot=axs[0, 1])
        axs[0, 1].set_title('Quantile-Quantile Plot')
        axs[0, 1].get_lines()[0].set_markerfacecolor('steelblue')
        axs[0, 1].get_lines()[1].set_color('red')

        # 3. CDF Comparison
        ecdf = np.sort(data), np.arange(1, len(data) + 1) / len(data)
        axs[1, 0].plot(*ecdf, 'b-', lw=2, label='Empirical CDF')
        axs[1, 0].plot(x, cdf, 'r--', lw=2, label='Theoretical CDF')
        axs[1, 0].set_title('Cumulative Distribution Comparison')
        axs[1, 0].legend()

        # 4. Boxplot with Theoretical Range
        axs[1, 1].boxplot(data, vert=False, patch_artist=True,
                          boxprops=dict(facecolor='steelblue', alpha=0.6))
        axs[1, 1].axvline(a, color='green', linestyle='--', label='Optimistic (a)')
        axs[1, 1].axvline(b, color='orange', linestyle='--', label='Most Likely (b)')
        axs[1, 1].axvline(c, color='red', linestyle='--', label='Pessimistic (c)')
        axs[1, 1].set_title('Data Range vs PERT Estimates')
        axs[1, 1].legend()

        for ax in axs.flat:
            ax.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()

    def generate_samples(self, group_id, max_attempts=5):
        """Generate samples with retry logic for KS test failures"""
        group = self.params_df[self.params_df['group_id'] == group_id].iloc[0]
        self.current_group = group

        # Parse parameters
        special_rows = list(map(int, group['special_rows'].split(',')))
        sample_size = group['sample_size']
        random_seed = group['random_seed']

        # Define parameters
        a_std, b_std, c_std = group['std_a'], group['std_b'], group['std_c']
        a_spec, b_spec, c_spec = group['spec_a'], group['spec_b'], group['spec_c']

        # Generate parameter sequence
        params = [(a_spec, b_spec, c_spec)] * 2 + [(a_std, b_std, c_std)] * 15 + [(a_spec, b_spec, c_spec)]

        attempt = 0
        while attempt < max_attempts:
            attempt += 1
            print(f"\n=== Group {group_id} Attempt {attempt}/{max_attempts} ===")

            # Generate samples with new random seed
            rng = np.random.default_rng(seed=random_seed + attempt)
            samples = np.array([
                self.pert_samples(a, b, c, size=sample_size, rng=rng)
                for a, b, c in params
            ])

            # Clip samples
            if group['clip_lower'] and group['clip_upper']:
                samples = self.clip_by_row(
                    samples, special_rows, a_spec, c_spec, a_std, c_std
                )

            self.samples = samples

            # Validate all rows
            if self.validate_all_rows():
                self.valid_samples[group_id] = samples
                return samples

        print(f"⚠️ Failed to generate valid samples for Group {group_id} after {max_attempts} attempts")
        return None

    def validate_all_rows(self):
        """Validate all rows in current sample set"""
        if self.samples is None:
            raise ValueError("No samples generated. Call generate_samples() first.")

        group = self.current_group
        a_std, b_std, c_std = group['std_a'], group['std_b'], group['std_c']
        a_spec, b_spec, c_spec = group['spec_a'], group['spec_b'], group['spec_c']

        params_all = [(a_spec, b_spec, c_spec)] * 2 + [(a_std, b_std, c_std)] * 15 + [(a_spec, b_spec, c_spec)]

        self.results = []
        for i, (a, b, c) in enumerate(params_all):
            res = self.check_pert_distribution(
                self.samples, a, b, c, group['group_id'], row_idx=i, plot=False
            )
            self.results.append(res)

        # Check if all rows pass KS test
        all_pass = all(r['ks_pass'] for r in self.results)
        print(f"\n=== Group {group['group_id']} Validation Summary ===")
        print(f"Sample size: {group['sample_size']}, Random seed: {group['random_seed']}")
        print(f"All rows pass KS test: {'✅ YES' if all_pass else '❌ NO'}")

        if all_pass:
            self._save_valid_samples()

        return all_pass

    def _save_valid_samples(self):
        """Save valid samples to CSV with metadata"""
        if self.samples is None:
            return

        group = self.current_group
        filename = f"valid_samples_{group['sample_filename']}.csv"

        # Create DataFrame with samples and metadata
        df_samples = pd.DataFrame(self.samples.T)

        # Add row type information
        special_rows = list(map(int, group['special_rows'].split(',')))
        row_types = []
        for i in range(len(self.samples)):
            if i in special_rows:
                row_types.append('special')
            else:
                row_types.append('standard')

        df_samples.columns = [f"{typ}_row_{i}" for i, typ in enumerate(row_types)]

        # Add metadata as first row
        metadata = {
            'group_id': group['group_id'],
            'sample_size': group['sample_size'],
            'random_seed': group['random_seed'],
            'ks_pass': True,
            'generation_time': pd.Timestamp.now().strftime('%Y/%m/%d %H:%M')
        }
        df_meta = pd.DataFrame([metadata])

        # Combine metadata and samples
        df_combined = pd.concat([df_meta, df_samples], axis=1)

        # Save to CSV
        df_combined.to_csv(filename, index=False)
        print(f"\n✅ Valid samples saved to {filename}")

    def generate_all_valid_samples(self):
        """Generate all valid sample groups from parameters"""
        for group_id in self.params_df['group_id']:
            self.generate_samples(group_id)

        print("\n=== Final Validation Summary ===")
        print(f"Total parameter groups: {len(self.params_df)}")
        print(f"Successfully generated: {len(self.valid_samples)}")
        print(f"Failed groups: {set(self.params_df['group_id']) - set(self.valid_samples.keys())}")

        return self.valid_samples

    def interactive_validation(self):
        """Interactive mode for exploring valid sample groups"""
        if not self.valid_samples:
            print("No valid samples available. Running generation first...")
            self.generate_all_valid_samples()

        print("\nAvailable valid sample groups:")
        valid_groups = pd.DataFrame([
            {
                'group_id': gid,
                'filename': self.params_df[self.params_df['group_id'] == gid]['sample_filename'].iloc[0],
                'sample_size': self.params_df[self.params_df['group_id'] == gid]['sample_size'].iloc[0]
            }
            for gid in self.valid_samples.keys()
        ])
        print(valid_groups)

        while True:
            try:
                group_id = int(input("\nEnter group ID to inspect (0 to exit): "))
                if group_id == 0:
                    break

                if group_id not in self.valid_samples:
                    print("Invalid group ID or group failed validation. Please try again.")
                    continue

                # Load the validated samples
                self.current_group = self.params_df[self.params_df['group_id'] == group_id].iloc[0]
                self.samples = self.valid_samples[group_id]

                # Allow viewing specific rows
                while True:
                    row_idx = input(
                        f"Enter row index (0-17) to view details for Group {group_id} "
                        f"('a' for all, 'n' for next group): "
                    )

                    if row_idx.lower() == 'n':
                        break
                    elif row_idx.lower() == 'a':
                        params_all = [
                                         (self.current_group['spec_a'], self.current_group['spec_b'],
                                          self.current_group['spec_c'])
                                     ] * 2 + [
                                         (self.current_group['std_a'], self.current_group['std_b'],
                                          self.current_group['std_c'])
                                     ] * 15 + [
                                         (self.current_group['spec_a'], self.current_group['spec_b'],
                                          self.current_group['spec_c'])
                                     ]
                        for i in range(len(params_all)):
                            a, b, c = params_all[i]
                            self.check_pert_distribution(
                                self.samples, a, b, c, group_id, row_idx=i, plot=True
                            )
                    elif row_idx.isdigit() and 0 <= int(row_idx) < 18:
                        i = int(row_idx)
                        if i in [0, 1, 17]:  # Special rows
                            a, b, c = (
                                self.current_group['spec_a'],
                                self.current_group['spec_b'],
                                self.current_group['spec_c']
                            )
                        else:  # Standard rows
                            a, b, c = (
                                self.current_group['std_a'],
                                self.current_group['std_b'],
                                self.current_group['std_c']
                            )
                        self.check_pert_distribution(
                            self.samples, a, b, c, group_id, row_idx=i, plot=True
                        )
                    else:
                        print("Invalid input. Please try again.")

            except ValueError:
                print("Please enter a valid number.")
            except Exception as e:
                print(f"An error occurred: {str(e)}")


# Main execution
if __name__ == "__main__":
    generator = PERTGenerator()

    # First generate all valid samples
    generator.generate_all_valid_samples()

    # Then enter interactive mode to inspect results
    generator.interactive_validation()